import os
import sys
from pathlib import Path
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
import logging
import torch
import gc
from tqdm import tqdm

class PDFProcessor:
    def __init__(self):
        # 配置日志
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
        
        # 设置环境变量
        os.environ['TORCH_DEVICE'] = 'cuda' if torch.cuda.is_available() else 'cpu'
        os.environ['BATCH_SIZE'] = '4'
        
        # 获取项目根目录
        self.root_dir = Path(__file__).parent.parent  # 获取参赛目录的路径
        
        # 设置输入输出目录（使用绝对路径）
        self.input_dir = self.root_dir / "rag" / "pdf"
        self.output_dir = self.root_dir / "rag" / "markdown"
        
        logging.info(f"PDF目录: {self.input_dir}")
        logging.info(f"Markdown目录: {self.output_dir}")
        
    def clean_gpu_memory(self):
        """清理GPU内存"""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            gc.collect()
            
    def save_image(self, img_data, img_name, img_dir):
        """保存图片并返回markdown链接"""
        try:
            # 如果img_data是PIL.Image对象，先转换为bytes
            if hasattr(img_data, 'save'):
                import io
                img_buffer = io.BytesIO()
                img_data.save(img_buffer, format='PNG')
                img_data = img_buffer.getvalue()
            
            img_path = img_dir / f"{img_name}.png"
            img_path.write_bytes(img_data)
            return f"![{img_name}](images/{img_path.name})"
        except Exception as e:
            logging.error(f"保存图片失败 {img_name}: {str(e)}")
            return ""  # 返回空字符串，这样即使图片保存失败也不会影响文本内容的保存
        
    def convert_pdfs(self):
        """批量转换PDF文件"""
        try:
            self.output_dir.mkdir(parents=True, exist_ok=True)
            
            # 加载模型
            logging.info("加载模型中...")
            artifact_dict = create_model_dict()
            converter = PdfConverter(artifact_dict=artifact_dict)
            logging.info("模型加载完成")
            
            # 获取所有PDF文件
            pdf_files = list(self.input_dir.rglob('*.pdf'))
            total_files = len(pdf_files)
            logging.info(f"找到 {total_files} 个PDF文件")
            
            # 转换处理
            success_count = 0
            for i, pdf_path in enumerate(tqdm(pdf_files, desc="转换进度")):
                try:
                    # 转换PDF
                    rendered = converter(str(pdf_path))
                    text, _, images = text_from_rendered(rendered)
                    
                    # 创建输出路径
                    rel_path = pdf_path.relative_to(self.input_dir)
                    output_path = self.output_dir / rel_path.with_suffix('.md')
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # 处理图片
                    if images:
                        img_dir = output_path.parent / 'images' / output_path.stem
                        img_dir.mkdir(parents=True, exist_ok=True)
                        
                        text += "\n\n## 图片\n\n"
                        for img_name, img_data in images.items():
                            img_link = self.save_image(img_data, img_name, img_dir)
                            text += f"{img_link}\n\n"
                    
                    # 保存Markdown
                    output_path.write_text(text, encoding='utf-8')
                    success_count += 1
                    
                    # 显示GPU内存使用情况
                    if torch.cuda.is_available():
                        allocated = torch.cuda.memory_allocated() / 1024**2
                        reserved = torch.cuda.memory_reserved() / 1024**2
                        logging.info(f"GPU内存: 已分配={allocated:.1f}MB, 保留={reserved:.1f}MB")
                    
                except Exception as e:
                    logging.error(f"转换失败 {pdf_path}: {str(e)}")
                    continue
                
                # 定期清理内存
                if (i + 1) % 10 == 0:
                    self.clean_gpu_memory()
            
            logging.info(f"转换完成! 成功: {success_count}/{total_files}")
            
        except Exception as e:
            logging.error(f"执行出错: {str(e)}")
        finally:
            self.clean_gpu_memory() 